package com.suntoon.framework.component;

import com.google.common.collect.ImmutableList;
import com.suntoon.framework.annotation.Limit;
import com.suntoon.framework.exception.MessageCenterException;
import com.suntoon.framework.util.LimitTypeEnum;
import org.apache.commons.lang3.StringUtils;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;

/**
 * @Author: LEAVES
 * @Version 1.0
 * @Date: 2021年06月21日 15时40分15秒
 * @Description:    限流处理aop
 */
@Aspect
@Component
public class LimitHandleAspect {
    @Resource
    private RedisTemplate<String, Object> redisTemplate;
    @Value("${limit.count}")
    private Integer LIMIT_COUNT;
    @Value("${limit.period}")
    public Integer LIMIT_PERIOD;

    //此处基于自身项目的路径做具体的设置
    @Pointcut("@within(com.suntoon.framework.annotation.Limit) || @annotation(com.suntoon.framework.annotation.Limit)")
    public void pointLimit() {
    }

    @Before("pointLimit()")
    public void limit(JoinPoint joinPoint) {
        ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        if (attributes == null) {
            return;
        }
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        Limit limitAnnotation = method.getAnnotation(Limit.class);
        if (limitAnnotation == null) {
            limitAnnotation = method.getDeclaringClass().getAnnotation(Limit.class);
        }
        LimitTypeEnum limitType = limitAnnotation.limitType();
        String key;
        int limitPeriod = limitAnnotation.period();
        if (limitPeriod == 1 && LIMIT_PERIOD != null) {
            limitPeriod = LIMIT_PERIOD;
        }
        int limitCount = limitAnnotation.count();
        if (limitCount == 5 && LIMIT_COUNT != null) {
            limitCount = LIMIT_COUNT;
        }
        switch (limitType) {
            case IP:
                key = getIpAddress(attributes);
                break;
            case CUSTOMER:
                key = limitAnnotation.keyGenerator();
                break;
            default:
                key = StringUtils.upperCase(method.getName());
        }
        ImmutableList<String> keys = ImmutableList.of(StringUtils.join(limitAnnotation.prefix(), key));
        String luaScript = buildLuaScript();
        RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
        Number count = redisTemplate.execute(redisScript, keys, limitCount, limitPeriod);
        if (count == null || count.intValue() > limitCount) {
            throw new MessageCenterException("访问过于繁忙！");
        }
    }

    /**
     * 限流 脚本
     *
     * @return lua脚本
     */
    public String buildLuaScript() {
        return "local c" +
                "\nc = redis.call('get',KEYS[1])" +
                // 调用不超过最大值，则直接返回
                "\nif c and tonumber(c) > tonumber(ARGV[1]) then" +
                "\nreturn c;" +
                "\nend" +
                // 执行计算器自加
                "\nc = redis.call('incr',KEYS[1])" +
                "\nif tonumber(c) == 1 then" +
                // 从第一次调用开始限流，设置对应键值的过期
                "\nredis.call('expire',KEYS[1],ARGV[2])" +
                "\nend" +
                "\nreturn c;";
    }

    private static final String UNKNOWN = "unknown";

    private String getIpAddress(ServletRequestAttributes attributes) {
        HttpServletRequest request = attributes.getRequest();
        String Xip = request.getHeader("X-Real-IP");
        String ip = request.getHeader("X-Forwarded-For");
        if(StringUtils.isNotEmpty(ip) && !UNKNOWN.equalsIgnoreCase(ip)){
            //多次反向代理后会有多个ip值，第一个ip才是真实ip
            int index = ip.indexOf(",");
            if(index != -1){
                return ip.substring(0,index);
            }else{
                return ip;
            }
        }
        ip = Xip;
        if(StringUtils.isNotEmpty(ip) && !UNKNOWN.equalsIgnoreCase(ip)){
            return ip;
        }
        if (StringUtils.isBlank(ip) || UNKNOWN.equalsIgnoreCase(ip)) {
            ip = request.getHeader("Proxy-Client-IP");
        }
        if (StringUtils.isBlank(ip) || UNKNOWN.equalsIgnoreCase(ip)) {
            ip = request.getHeader("WL-Proxy-Client-IP");
        }
        if (StringUtils.isBlank(ip) || UNKNOWN.equalsIgnoreCase(ip)) {
            ip = request.getHeader("HTTP_CLIENT_IP");
        }
        if (StringUtils.isBlank(ip) || UNKNOWN.equalsIgnoreCase(ip)) {
            ip = request.getHeader("HTTP_X_FORWARDED_FOR");
        }
        if (StringUtils.isBlank(ip) || UNKNOWN.equalsIgnoreCase(ip)) {
            ip = request.getRemoteAddr();
        }
        return ip;
    }
}
